import numpy as np
import onnxruntime as ort
from PIL import Image
import torchvision.transforms as transforms

# 加载 ONNX 模型
model_path = 'mobilenet_v3_large.onnx'
ort_session = ort.InferenceSession(model_path)

# 加载标签名
def load_labels(label_file):
    with open(label_file, 'r') as f:
        labels = [line.strip() for line in f.readlines()]
    return labels

# 加载图片并进行预处理
def load_image(image_path):
    image = Image.open(image_path).convert('RGB')  # 确保为 RGB 格式
    preprocess = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.CenterCrop(224),  # MobileNetV3 期望输入 224x224 的图片
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
    ])
    image_tensor = preprocess(image)
    return image_tensor.unsqueeze(0).numpy()  # 增加一个批次维度

# 进行推理
def infer(image_path, labels):
    input_image = load_image(image_path)
    ort_inputs = {ort_session.get_inputs()[0].name: input_image}
    ort_outs = ort_session.run(None, ort_inputs)

    # 获取预测结果
    probabilities = ort_outs[0]
    top5_indices = np.argsort(probabilities)[0, -5:][::-1]  # 获取前5个类的索引
    top5_probabilities = probabilities[0, top5_indices]  # 获取前5个类的概率

    return top5_indices, top5_probabilities

# 指定图片路径和标签文件路径
image_path = '22.png'
label_file = 'names.txt'  # 确保这个文件中包含每个类的名称
labels = load_labels(label_file)

top5_indices, top5_probabilities = infer(image_path, labels)

# 打印预测结果
print('Top 5 predicted classes:')
for i in range(5):
    print(f'{labels[top5_indices[i]]}: {top5_probabilities[i]:.4f}')
